#!/usr/bin/env python
# -*- coding:utf-8 -*-
# Create by zhang
# Create on 2022/6/22 22:01
import datetime

import pandas as pd
import os
from enum import Enum
import time
from config.config import config
from core.enums import Period, DateFormat
from domain.transaction_data.spider import tushare_spider, xueqiu_spider
from infrastructure.util import dateutils
from core.dataClasses import StockTradeDataColumnName


def cvs_reader_only(path):
    assert os.path.exists(path)
    return pd.read_csv(path)


def cvs_reader(path, index_col='date_index', parse_dates=[StockTradeDataColumnName.DATE_INDEX.value]) -> pd.DataFrame:
    return pd.read_csv(path, index_col=index_col, parse_dates=parse_dates)


def cvs_write(path, df: pd.DataFrame):
    df.to_csv(path_or_buf=path)


def get_stock_transaction_data(ts_code, period:Period = Period.day, endTime: datetime.datetime = None) -> pd.DataFrame:
    return _fetchAndUpdateKlineByXueqiu(ts_code, period, endTime)


def realtimeQuote(security: str) -> pd.Series:
    return xueqiu_spider.realtimeQuote(security)


def _get_transaction_data_by_tushare(ts_code, download=True, update=False) -> pd.DataFrame:
    """
    获取包含最新数据的股票交易数据，以日为单位
    :param ts_code:
    :return:
    """
    path = config.get('data', 'folder') + os.sep + "day" + os.sep + ts_code + ".csv"
    try:
        df = cvs_reader_only(path).sort_index(ascending=True)  # type:pd.DataFrame
        date = df.index.tolist()[-1]
    except:
        date = 0
        df = pd.DataFrame()
    if not df.empty and not update:
        return df
    if df.empty and not download:
        return df
    to_update = False
    lastest_transaction_dateStr = dateutils.getLastestTradingDate(time.strftime(DateFormat.Y_m_d.value)).replace('-',
                                                                                                                  '')
    laster_transaction_dateStr = dateutils.getLastestTradingDate(
        dateutils.get_dateStr_before_num_days_by_dateStr(1, lastest_transaction_dateStr, dateFormat=DateFormat.Ymd), dateFormat=DateFormat.Ymd)
    if date < int(lastest_transaction_dateStr) and (int(time.strftime('%Y%m%d%H%M%S')) > int(
            lastest_transaction_dateStr + '150000') or date <= int(laster_transaction_dateStr)):
        df = tushare_spider.daily_data(ts_code=ts_code)  # type:pd.DataFrame
        to_update = True
    df['date_index'] = pd.to_datetime(df['trade_date'], format=DateFormat.Ymd.value)
    df.set_index(keys='date_index', inplace=True)
    df = df.sort_index(ascending=True)
    if not df.empty and to_update:
        df.to_csv(path)
    return df


def _fetchAndUpdateKlineByXueqiu(ts_code, period:Period = Period.day, endTime: datetime.datetime = None) -> pd.DataFrame:
    """
    获取包含最新数据的股票交易数据
    :param ts_code:
    :param period: 股票周期
    :param endTime: 股票最后日期
    :return:
    """
    subPath = Period.day.name
    if period == Period.minute30:
        subPath = Period.minute30.name
    elif period == Period.minute5:
        subPath = Period.minute5.name
    elif period == Period.minute1:
        subPath = Period.minute1.name
    elif period == Period.week:
        subPath = Period.week.name
    path = config.get('data', 'folder') + os.sep + subPath + os.sep + ts_code + ".csv"
    try:
        df:pd.DataFrame = pd.read_csv(path, parse_dates=[StockTradeDataColumnName.DATE_INDEX.value,], infer_datetime_format=True)

        # df = cvs_reader_only(path).sort_index(ascending=True)  # type:pd.DataFrame
        df.set_index(keys=StockTradeDataColumnName.DATE_INDEX.value, inplace=True)
        date = df[StockTradeDataColumnName.TRADE_DATE.value].tolist()[-1]
    except:
        date = 0
        df = pd.DataFrame()
    if df.empty:
        df = xueqiu_spider.kline(ts_code, period, days=10000)  # type:pd.DataFrame
        if not df.empty:
            df.to_csv(path)
    else:
        lastest_transaction_datetimeStr = dateutils.getPeriodLastestTradingDatetime(period).strftime(DateFormat.YmdHMS.value)
        if date < int(lastest_transaction_datetimeStr):
            dt = datetime.datetime.strptime(str(int(date / 1000000)), DateFormat.Ymd.value)
            count = dateutils.countTradeDays(begin=datetime.date(dt.year, dt.month, dt.day)) + 2
            multiple = 1
            if period == Period.minute30:
                multiple = 4 * 2
            elif period == Period.minute5:
                multiple = 4 * 2 * 6
            elif period == Period.minute1:
                multiple = 4 * 2 * 6 * 1
            newer_df = xueqiu_spider.kline(ts_code, period, days=count*multiple)  # type:pd.DataFrame
            if not newer_df.empty:
                df = pd.concat([df, newer_df], axis=0)
                df = df[~df.index.duplicated(keep='last')]
                df.sort_index(ascending=True)
                df.to_csv(path)
    if endTime is not None:
        endDate = endTime.strftime(DateFormat.Y_m_d_H_M_S.value) if period == Period.minute30 else endTime.strftime(DateFormat.Y_m_d.value)
        df = df.loc[:endDate]
    return df